"""DHF模型产流计算模块

实现大伙房模型的产流计算，包括不透水面积径流、地面壤中流和地下径流的计算。
"""
import math
from typing import Tuple, Dict

import numpy as np

from .parameters import DHFParameters


class RunoffGeneration:
    """DHF模型产流计算类"""

    def __init__(self, parameters: DHFParameters):
        """
        初始化产流计算模块
        
        Args:
            parameters: DHF模型参数对象
        """
        self.params = parameters
        self.runoff_params = parameters.runoff_params

        # 验证和修正关键参数
        if 'B' in self.runoff_params:
            self.runoff_params['B'] = max(1.01, min(5.0, self.runoff_params['B']))
        if 'K' in self.runoff_params:
            self.runoff_params['K'] = max(0.01, min(2.0, self.runoff_params['K']))
        if 'U0' in self.runoff_params:
            self.runoff_params['U0'] = max(1.0, self.runoff_params['U0'])
        if 'S0' in self.runoff_params:
            self.runoff_params['S0'] = max(1.0, self.runoff_params['S0'])
        if 'D0' in self.runoff_params and 'U0' in self.runoff_params:
            self.runoff_params['D0'] = max(self.runoff_params['U0'], self.runoff_params['D0'])

        # 初始化状态变量
        self.S_a = self.runoff_params['S0'] * 0.5  # 表层蓄水量初值
        self.U_a = self.runoff_params['U0'] * 0.5  # 下层蓄水量初值
        self.V_a = self.runoff_params['V0'] * 0.5  # 地下水库蓄水量初值

        # 计算D_a
        self.D_a = self.U_a + self.V_a

    def calculate_impervious_runoff(self, P: float, E_D: float) -> float:
        """
        计算不透水面积上的径流 (公式1-2)
        
        Args:
            P: 时段降雨量 (mm)
            E_D: 流域蒸发能力 (mm)
            
        Returns:
            不透水面积径流 y_c (mm)
        """
        # 计算净降雨量 P_E (公式2)
        K_C = self.params.evaporation_params['K_C']
        C = 24  # 日降雨观测的时段数（小时）
        P_E = P - K_C * E_D / C

        # 计算不透水面积径流 (公式1)
        g = self.runoff_params['g']
        y_c = g * max(0, P_E)

        return y_c

    def calculate_surface_deduction(self, P_c: float) -> float:
        """
        计算表层扣损 R (公式4-8)
        
        Args:
            P_c: 净降雨量 (mm)
            
        Returns:
            表层扣损 R (mm)
        """
        S0 = self.runoff_params['S0']
        A = self.runoff_params['A']

        # 计算S_m (公式6)
        if self.S_a <= 0:
            S_m = 0
        else:
            S_m = A * S0 * (1 - (1 - self.S_a / S0) ** (1 / A))

        # 计算表层最大蓄水量 S_M (公式8)
        S_M = A * S0

        # 计算表层扣损 R
        if S_m + P_c < A * S0:
            # 公式4
            if S_m + P_c > 0:
                R = P_c + self.S_a - S0 * (1 - (S_m + P_c) / (A * S0)) ** A
            else:
                R = 0
        else:
            # 公式5
            R = P_c - (S0 - self.S_a)

        return max(0, R)

    def calculate_infiltration_coefficients(self, t: float = 1.0) -> Tuple[float, float]:
        """
        计算下渗系数 Z1 和 Z2
        
        Args:
            t: 时段长度 (天)
            
        Returns:
            (Z1, Z2) 下渗系数元组
        """
        K = self.runoff_params['K']
        U0 = self.runoff_params['U0']
        D0 = self.runoff_params['D0']

        # 计算k1和k2 (公式12)
        k2 = K
        k1 = k2 * U0 / D0

        # 计算Z1和Z2
        Z1 = 1 - np.exp(-k1 * U0 / D0)
        Z2 = 1 - np.exp(-k2 * t)

        return Z1, Z2

    def calculate_D_a_from_U_a(self, U_a: float) -> float:
        """
        根据U_a计算D_a (公式13)
        
        Args:
            U_a: 下层蓄水量 (mm)
            
        Returns:
            D_a: 下层蓄水量与地下水库蓄水量之和 (mm)
        """
        U0 = self.runoff_params['U0']
        D0 = self.runoff_params['D0']

        if U_a <= 0:
            return 0

        D_a = D0 * (1 - (1 - U_a / U0) ** (U0 / D0))
        return D_a

    def calculate_D_n_and_U_n(self, U_a: float) -> Tuple[float, float]:
        """
        计算D_n和U_n (公式18-19)
        
        Args:
            U_a: 下层蓄水量 (mm)
            
        Returns:
            (D_n, U_n) 元组
        """
        B = self.runoff_params['B']
        U0 = self.runoff_params['U0']
        D0 = self.runoff_params['D0']

        if U_a <= 0:
            return 0, 0

        # 公式18
        D_n = B * D0 * (1 - (1 - U_a / U0) ** (U0 / (B * D0)))

        # 公式19
        U_n = B * U0 * (1 - (1 - U_a / U0) ** (1 / B))

        return D_n, U_n

    def calculate_total_runoff(self, R: float, U_a: float, t: float = 1.0) -> float:
        """
        计算透水面积总径流 y (公式14-15)
        """
        if R <= 0:
            return 0.0
        
        # 获取参数
        B = self.runoff_params.get('B', 1.2)
        U0 = self.runoff_params.get('U0', 100.0)
        K = self.runoff_params.get('K', 0.3)
        
        # 确保参数合理性
        B = max(1.01, B)
        U0 = max(1.0, U0)
        U_a = max(0.0, min(U_a, U0))
        
        # 计算Z2
        Z2 = 1 - math.exp(-K * t)
        Z2 = max(0.001, min(0.999, Z2))
        
        # 计算u_n
        if U_a >= U0:
            u_n = B * U0
        else:
            ratio = U_a / U0
            u_n = B * U0 * (1 - (1 - ratio) ** (1/B))
        
        # 判断条件
        condition_value = R + Z2 * u_n
        threshold = Z2 * B * U0
        
        if condition_value >= threshold:
            # 公式15
            y = R - Z2 * (U0 - U_a)
        else:
            # 公式14 - 添加数值稳定性检查
            try:
                ratio = condition_value / threshold
                if ratio <= 0:
                    y = R + Z2 * (U_a - U0)
                else:
                    exponent = -B
                    base = 1 - ratio
                    if base <= 0:
                        y = R - Z2 * (U0 - U_a)
                    else:
                        power_term = base ** exponent
                        if math.isinf(power_term) or math.isnan(power_term):
                            y = R - Z2 * (U0 - U_a)
                        else:
                            y = R + Z2 * (U_a - U0) + Z2 * U0 * power_term
            except (OverflowError, ZeroDivisionError, ValueError):
                y = R - Z2 * (U0 - U_a)
        
        # 确保结果非负且合理
        y = max(0.0, y)
        if y > R * 10:  # 防止异常大值
            y = R
        
        return y
    
    def calculate_interflow(self, R: float, U_a: float, t: float = 1.0) -> float:
        """
        计算地面壤中流 y_u (公式16-17)
        """
        if R <= 0:
            return 0.0
    
        # 获取参数并确保合理性
        B = max(1.01, self.runoff_params.get('B', 1.2))
        U0 = max(1.0, self.runoff_params.get('U0', 100.0))
        D0 = max(U0, self.runoff_params.get('D0', 130.0))
        K = self.runoff_params.get('K', 0.3)
    
        U_a = max(0.0, min(U_a, U0))
    
        # 计算Z1
        k1 = K * U0 / D0
        Z1 = 1 - math.exp(-k1 * t)
        Z1 = max(0.001, min(0.999, Z1))
    
        # 计算D_n
        if U_a >= U0:
            D_n = B * D0
        else:
            ratio = U_a / U0
            exponent = U0 / (B * D0)
            D_n = B * D0 * (1 - (1 - ratio) ** exponent)
    
        # 判断条件
        condition_value = Z1 * D_n + R
        threshold = Z1 * B * D0
    
        # 基础项
        base_term = Z1 * D0 * (1 - U_a/U0) ** (U0/D0)
    
        if condition_value >= threshold:
            # 公式17
            y_u = R - base_term
        else:
            # 公式16 - 添加数值稳定性
            try:
                ratio = condition_value / threshold
                if ratio <= 0 or ratio >= 1:
                    y_u = R - base_term
                else:
                    power_term = (1 - ratio) ** B
                    if math.isinf(power_term) or math.isnan(power_term):
                        y_u = R - base_term
                    else:
                        y_u = R - base_term + Z1 * D0 * power_term
            except (OverflowError, ZeroDivisionError, ValueError):
                y_u = R - base_term
    
        # 确保结果合理
        y_u = max(0.0, min(y_u, R))
    
        return y_u

    def calculate_groundwater_runoff(self, y: float, y_u: float) -> float:
        """
        计算地下径流 y_L (公式20-21)
        
        Args:
            y: 总径流 (mm)
            y_u: 地面壤中流 (mm)
            
        Returns:
            地下径流 y_L (mm)
        """
        K_W = self.runoff_params['K_W']

        # 计算地下水库下渗强度 R_L (公式21)
        R_L = y - y_u

        # 计算地下径流 (公式20)
        y_L = K_W * max(0, R_L)

        return y_L

    def update_surface_storage(self, P_c: float, R: float, t: float = 1.0):
        """
        更新表层蓄水量 (公式22-24)
        
        Args:
            P_c: 净降雨量 (mm)
            R: 表层扣损 (mm)
            t: 时段长度 (天)
        """
        S0 = self.runoff_params['S0']
        A = self.runoff_params['A']

        if P_c > 0:  # 有降雨时
            # 计算S_m
            if self.S_a <= 0:
                S_m = 0
            else:
                S_m = A * S0 * (1 - (1 - self.S_a / S0) ** (1 / A))

            if S_m + P_c >= A * S0:
                # 公式23
                self.S_a = S0
            else:
                # 公式22
                if S_m + P_c > 0:
                    self.S_a = S0 * (1 - (1 - (S_m + P_c) / (A * S0)) ** A)
                else:
                    self.S_a = max(0, self.S_a + P_c - R)
        else:  # 蒸发时
            # 公式30
            E_u = min(abs(P_c), self.S_a)  # 表层可蒸发量
            self.S_a = max(0, self.S_a - E_u)

    def update_subsurface_storage(self, R: float, y: float, t: float = 1.0):
        """
        更新下层蓄水量 (公式25-27)
        
        Args:
            R: 表层扣损 (mm)
            y: 总径流 (mm)
            t: 时段长度 (天)
        """
        U0 = self.runoff_params['U0']
        B = self.runoff_params['B']

        Z1, Z2 = self.calculate_infiltration_coefficients(t)
        D_n, U_n = self.calculate_D_n_and_U_n(self.U_a)

        if R > 0:  # 有入渗时
            if Z2 * U_n + R >= Z2 * B * U0:
                # 公式27
                self.U_a = self.U_a + Z2 * (U0 - self.U_a)
            else:
                # 公式26
                if Z2 * U_n + R > 0:
                    self.U_a = self.U_a + Z2 * (U0 - self.U_a) - Z2 * U0 * (1 - (Z2 * U_n + R) / (Z2 * B * U0)) ** B
                else:
                    self.U_a = self.U_a + Z2 * (U0 - self.U_a)
        else:  # 蒸发时
            # 公式33
            E_c = abs(R)  # 总蒸发需求
            E_u = min(E_c, self.S_a)  # 表层实际蒸发
            E_L = (E_c - E_u) * self.U_a / U0  # 下层蒸发
            self.U_a = max(0, self.U_a - E_L)

        # 确保不超过容量
        self.U_a = min(self.U_a, U0)

        # 更新D_a
        self.D_a = self.calculate_D_a_from_U_a(self.U_a)

    def calculate_runoff_components(self, P: float, E_D: float, t: float = 1.0) -> Dict[str, float]:
        """
        计算完整的产流过程
        
        Args:
            P: 时段降雨量 (mm)
            E_D: 流域蒸发能力 (mm)
            t: 时段长度 (天)
            
        Returns:
            包含各径流分量的字典
        """
        # 1. 计算不透水面积径流
        y_c = self.calculate_impervious_runoff(P, E_D)

        # 2. 计算净降雨量 P_c (公式9)
        P_c = P - E_D - y_c

        # 3. 计算表层扣损
        R = self.calculate_surface_deduction(P_c)

        # 4. 计算各径流分量
        y = self.calculate_total_runoff(R, self.U_a, t)
        y_u = self.calculate_interflow(R, self.U_a, t)
        y_L = self.calculate_groundwater_runoff(y, y_u)

        # 5. 更新状态变量
        self.update_surface_storage(P_c, R, t)
        self.update_subsurface_storage(R, y, t)

        # 6. 计算总径流 - 修复错误
        total_runoff = y_c + y  # ✅ 正确的计算方式

        return {
            'impervious_runoff': y_c,  # 不透水面积径流
            'interflow': y_u,  # 地面壤中流
            'groundwater_runoff': y_L,  # 地下径流
            'total_runoff': total_runoff,  # 总径流
            'surface_storage': self.S_a,  # 表层蓄水量
            'subsurface_storage': self.U_a,  # 下层蓄水量
            'net_rainfall': P_c,  # 净降雨量
            'surface_deduction': R  # 表层扣损
        }

    def get_state_variables(self) -> Dict[str, float]:
        """
        获取当前状态变量
        
        Returns:
            状态变量字典
        """
        return {
            'S_a': self.S_a,  # 表层蓄水量
            'U_a': self.U_a,  # 下层蓄水量
            'V_a': self.V_a,  # 地下水库蓄水量
            'D_a': self.D_a  # 下层蓄水量与地下水库蓄水量之和
        }

    def set_state_variables(self, S_a: float = None, U_a: float = None, V_a: float = None):
        """
        设置状态变量
        
        Args:
            S_a: 表层蓄水量 (mm)
            U_a: 下层蓄水量 (mm)
            V_a: 地下水库蓄水量 (mm)
        """
        if S_a is not None:
            self.S_a = max(0, min(S_a, self.runoff_params['S0']))
        if U_a is not None:
            self.U_a = max(0, min(U_a, self.runoff_params['U0']))
        if V_a is not None:
            self.V_a = max(0, min(V_a, self.runoff_params['V0']))

        # 更新D_a
        self.D_a = self.U_a + self.V_a
